package com.sunday.common.mq.rocket.brave;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.internal.Nullable;
import brave.messaging.MessagingRequest;
import brave.propagation.ThreadLocalCurrentTraceContext;
import brave.propagation.TraceContext;
import brave.propagation.TraceContext.Extractor;
import brave.propagation.TraceContext.Injector;
import brave.propagation.TraceContextOrSamplingFlags;
import brave.sampler.SamplerFunction;
import com.sunday.common.mq.rocket.brave.propagation.SpansAndScope;
import com.sunday.common.mq.rocket.brave.propagation.TracingRocketPostProcessor;
import lombok.extern.slf4j.Slf4j;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.rocketmq.client.hook.ConsumeMessageContext;
import org.apache.rocketmq.client.hook.ConsumeMessageHook;
import org.apache.rocketmq.client.trace.TraceConstants;
import org.apache.rocketmq.common.message.MessageConst;
import org.apache.rocketmq.common.message.MessageExt;

import java.util.Arrays;

import static brave.Span.Kind.CONSUMER;

/**
 * Pull brave 无法利用原始操作进行处理，需要加持 AOP
 *
 * @author zsy
 * @see brave.spring.rabbit.TracingRabbitListenerAdvice#invoke(MethodInvocation)
 * @see org.apache.rocketmq.client.trace.hook.ConsumeMessageOpenTracingHookImpl
 * @since 2023/7/31
 */
@Slf4j
final class PullConsumeMessageBraveTracingHookImpl implements ConsumeMessageHook {

    final RocketTracing rocketTracing;
    final Tracing tracing;
    final Tracer tracer;
    final Extractor<MessageConsumerRequest> extractor;
    final Injector<MessageConsumerRequest> injector;
    final SamplerFunction<MessagingRequest> sampler;
    @Nullable
    final String remoteServiceName;

    PullConsumeMessageBraveTracingHookImpl(RocketTracing rocketTracing) {
        this.rocketTracing = rocketTracing;
        this.tracing = rocketTracing.tracing;
        this.tracer = tracing.tracer();
        this.extractor = rocketTracing.consumerExtractor;
        this.sampler = rocketTracing.consumerSampler;
        this.injector = rocketTracing.consumerInjector;
        this.remoteServiceName = rocketTracing.remoteServiceName;
    }

    @Override
    public String hookName() {
        return "ConsumeMessageBraveTracingHook";
    }

    @Override
    public void consumeMessageBefore(ConsumeMessageContext context) {
        if (context == null || context.getMsgList() == null || context.getMsgList().isEmpty()) {
            return;
        }

        for (MessageExt msg : context.getMsgList()) {
            if (msg == null) {
                continue;
            }

            MessageConsumerRequest request = new MessageConsumerRequest(msg);

            TraceContextOrSamplingFlags extracted =
                    rocketTracing.extractAndClearTraceIdHeaders(extractor, request, msg);

            Span consumerSpan = rocketTracing.nextMessagingSpan(sampler, request, extracted);
            Span listenerSpan = tracer.newChild(consumerSpan.context());

            if (!consumerSpan.isNoop()) {
                setConsumerSpan(consumerSpan, msg);
                // 只产生一次时间戳开销
                long timestamp = tracing.clock(consumerSpan.context()).currentTimeMicroseconds();
                consumerSpan.start(timestamp);
                long consumerFinish = timestamp + 1L; // save a clock reading
                consumerSpan.finish(consumerFinish);
                // 没有使用作用域跨度，因为我们希望从预配置的时间开始
                listenerSpan.name("on-message").start(consumerFinish);
            }

            /**
             * Tracer.SpanInScope.close() 方法和 span.finish() 方法不完全等同，但它们通常用于相似的目的。
             * Tracer.SpanInScope.close() 方法是在使用某些分布式追踪系统或库时，在特定的作用域中关闭当前 span 的一种约定方式。该方法会结束当前作用域内的 span，并将其从活动的 span 列表中移除。
             * 而 span.finish() 方法是结束和完成一个具体的 span。它会记录 span 的结束时间，并将相关信息发送给追踪系统进行后续处理和可视化。
             * Tracer.SpanInScope ws = tracer.withSpanInScope(listenerSpan)
             * 1. Tracer.SpanInScope.close() 中使用 {@link ThreadLocalCurrentTraceContext#newScope(TraceContext)} 存在 local.set(currentSpan) 这个只能使用 close() 进行 local.set(null) 清除
             * 2. span.finish() 不会参与 ThreadLocal 的清理动作，只会进行链路的完结
             */
            TracingRocketPostProcessor.release(null);
            TracingRocketPostProcessor.set(new SpansAndScope(Arrays.asList(listenerSpan), tracer.withSpanInScope(listenerSpan)));
//            log.info("{}", extracted.context());
        }
//        context.setMqTraceContext(spanScopeList);
    }

    void setConsumerSpan(Span span, MessageExt msg) {
        span.kind(CONSUMER).name(TraceConstants.FROM_PREFIX + msg.getTopic());
        maybeTag(span, RocketTracing.ROCKETMQ_TOPIC, msg.getTopic());
        maybeTag(span, TraceConstants.ROCKETMQ_MSG_ID, msg.getMsgId());
        maybeTag(span, TraceConstants.ROCKETMQ_TAGS, msg.getTags());
        maybeTag(span, TraceConstants.ROCKETMQ_KEYS, msg.getKeys());
        maybeTag(span, TraceConstants.ROCKETMQ_BODY_LENGTH, String.valueOf(msg.getStoreSize()));
        maybeTag(span, TraceConstants.ROCKETMQ_RETRY_TIMERS, String.valueOf(msg.getReconsumeTimes()));
        maybeTag(span, TraceConstants.ROCKETMQ_REGION_ID, msg.getProperty(MessageConst.PROPERTY_MSG_REGION));
    }

    void maybeTag(Span span, String tag, String value) {
        if (value != null) span.tag(tag, value);
    }

    @Override
    public void consumeMessageAfter(ConsumeMessageContext context) {
//        if (context == null || context.getMsgList() == null || context.getMsgList().isEmpty()) {
//            return;
//        }
//        List<SpanScope> spanScopeList = (List<SpanScope>) context.getMqTraceContext();
//        if (spanScopeList == null) {
//            return;
//        }
//        for (SpanScope spanScope : spanScopeList) {
//            Span listenerSpan = spanScope.span();
//            Throwable error = null;
//            try {
//                maybeTag(listenerSpan, TraceConstants.ROCKETMQ_SUCCESS, String.valueOf(context.isSuccess()));
//            } catch (Throwable t) {
//                error = t;
//                log.error("{}", t.toString(), t);
//                throw t;
//            } finally {
//                if (error != null) listenerSpan.error(error);
//                // 因为 NoopSpan finish 什么都没有做，所以不需要判定.isNoop()，上面那么写是因为不行进行额外的其他操作
//                listenerSpan.finish();
//                spanScope.scope().close();
//            }
//        }
    }
}
